from llama_index.core import (
    StorageContext,
    load_index_from_storage,
    VectorStoreIndex,
    Document,
)
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore

# save to disk
# 构建示例文档内容列表
sample_texts = [
    "LlamaIndex支持持久化存储，可以保存和重用嵌入向量。",
    "通过持久化存储，可以避免重复计算嵌入向量，节省时间和成本。",
    "SimpleVectorStore是LlamaIndex提供的一个简单的向量存储实现。",
]
# 导入OpenAIEmbedding嵌入模型
from llama_index.embeddings.openai import OpenAIEmbedding

# 设置嵌入模型为OpenAIEmbedding（默认使用text-embedding-ada-002）
embed_model = OpenAIEmbedding()
# 将每段文本包装为Document对象，便于后续索引处理
documents = [Document(text=text) for text in sample_texts]
db = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = db.get_or_create_collection("quickstart")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

index = VectorStoreIndex.from_documents(
    documents, storage_context=storage_context, embed_model=embed_model
)
